import halo
import time
from functools import wraps
import sys

def halo_decorator(text='Running...', spinner='dots', color='cyan', enabled=True):
    """
    一个装饰器，使用 halo 模块在函数运行时显示 spinner 动画。

    Args:
        text (str, optional): spinner 前面显示的消息. Defaults to 'Loading...'.
        spinner (str or dict, optional):  spinner 类型，可以是 halo 库预定义的类型名 (例如 'dots', 'bouncingBall')
                                          或者自定义的 spinner 帧和间隔字典. Defaults to 'dots'.
                                          参考 halo 文档: http://halo.josealerma.com/
        color (str, optional): spinner 颜色. Defaults to 'blue'.
                                 可选颜色: grey, red, green, yellow, blue, magenta, cyan, white
        enabled (bool, optional): 是否启用 spinner 动画. Defaults to True.
                                  如果设置为 False, 则不显示 spinner, 直接执行函数.

    Returns:
        function: 装饰器函数.
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not enabled: # 如果 disabled, 直接执行函数，不显示 spinner
                return func(*args, **kwargs)

            spin = halo.Halo(text=text, spinner=spinner, color=color, stream=sys.stderr) # stream=sys.stderr 避免输出到 stdout 干扰程序输出
            result = None
            try:
                spin.start() # 启动 spinner
                result = func(*args, **kwargs) # 执行被装饰的函数
            except Exception as e:
                spin.fail(text=f"{text} Failed! Error: {e}") # 函数执行失败，spinner 变为失败状态并显示错误信息
                raise e # 重新抛出异常，方便调用者处理
            else:
                spin.succeed(text=f"{text} Done!") # 函数成功执行完毕，spinner 变为成功状态并显示完成信息
            finally:
                spin.stop() # 确保 spinner 停止 (虽然 succeed/fail 也会停止，但显式 stop 更保险)
            return result # 返回被装饰函数的返回值
        return wrapper
    return decorator

if __name__ == '__main__':
    @halo_decorator(text='Processing files', spinner='earth', color='green')
    def process_files():
        """模拟文件处理"""
        time.sleep(2) # 模拟耗时操作
        return ["file1.txt", "file2.txt", "file3.txt"]

    @halo_decorator(text='Calculating results', spinner='clock', color='yellow')
    def calculate_results(data):
        """模拟计算结果"""
        time.sleep(1.5)
        return sum(data)

    @halo_decorator(text='Simulating error', spinner='monkey', color='magenta')
    def simulate_error():
        """模拟一个会抛出错误的函数"""
        time.sleep(1)
        raise ValueError("Oops! Something went wrong during simulation.")

    @halo_decorator(enabled=False) # 禁用 spinner
    def no_spinner_function():
        """一个没有 spinner 的函数"""
        time.sleep(4)
        print("Function without spinner executed.3")


    processed_files = process_files()
    print(f"Processed files: {processed_files}")

    calculation_result = calculate_results([10, 20, 30])
    print(f"Calculation result: {calculation_result}")

    try:
        simulate_error()
    except ValueError as e:
        print(f"Caught error: {e}")

    no_spinner_function() # 调用没有 spinner 的函数

# import sys
# import time
# # 将模块所在目录添加到 sys.path
# sys.path.append('/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/py/utils')
# from yk_tools import halo_decorator

# @halo_decorator.halo_decorator(text='Running')
# def foo():
#     time.sleep(2)
    
# foo()